import torch
import torch.nn as nn
from torch.nn.parameter import Parameter


class NonLocal_Module_Pos(nn.Module):
    """ NonLocal_module is similar to Position attention module of DANet (CVPR2019)"""
    #Ref from SAGAN
    def __init__(self, in_dim):
        super(NonLocal_Module_Pos, self).__init__()
        self.chanel_in = in_dim

        self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim//2, kernel_size=1)
        self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim//2, kernel_size=1)
        self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)
        self.gamma = Parameter(torch.zeros(1))

        self.softmax = nn.Softmax(dim=-1)
    def forward(self, x):
        """
            inputs :
                x : input feature maps( B X C X H X W)
            returns :
                out : attention value + input feature
                attention: B X (HxW) X (HxW)
        """
        m_batchsize, C, height, width = x.size()
        proj_query = self.query_conv(x).view(m_batchsize, -1, width*height).permute(0, 2, 1)
        proj_key = self.key_conv(x).view(m_batchsize, -1, width*height)
        energy = torch.bmm(proj_query, proj_key)
        attention = self.softmax(energy)
        proj_value = self.value_conv(x).view(m_batchsize, -1, width*height)

        out = torch.bmm(proj_value, attention.permute(0, 2, 1))
        out = out.view(m_batchsize, C, height, width)

        out = self.gamma*out + x
        return out


class NonLocal_Module(nn.Module):
    """ NonLocal_module is similar to Position attention module of DANet (CVPR2019)"""
    #Ref from SAGAN
    def __init__(self, seqLen, in_dim, reduce_ration=8):
        super(NonLocal_Module, self).__init__()
        self.seqLen = seqLen
        self.chanel_in = in_dim
        self.avgpool = nn.AvgPool3d((seqLen, 1, 1), stride=1)
        self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim//reduce_ration, kernel_size=1)
        self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim//reduce_ration, kernel_size=1)
        self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)
        self.gamma = Parameter(torch.zeros(1))

        self.softmax = nn.Softmax(dim=-1)
    def forward(self, x):
        """
            inputs :
                x : input feature maps( B X C X T X H X W)
            returns :
                out : attention value + input feature
                attention: B X (HxW) X (HxW)
        """
        m_batchsize, C, temp, height, width = x.size()
        assert temp == self.seqLen

        x_pool = self.avgpool(x).squeeze(2)
        # proj_query (b, WxH, C)
        # proj_key (b, C, WxH)
        proj_query = self.query_conv(x_pool).view(m_batchsize, -1, width*height).permute(0, 2, 1)
        proj_key = self.key_conv(x_pool).view(m_batchsize, -1, width*height)
        # energy (b, WxH, WxH)
        energy = torch.bmm(proj_query, proj_key)
        # attention (b, WxH, WxH)
        # proj_value (b, C, WxH)
        attention = self.softmax(energy)
        proj_value = self.value_conv(x_pool).view(m_batchsize, -1, width*height)
        # out = (b, C, WxH) x (b, WxH, WxH) = (b, C, WxH)
        out = torch.bmm(proj_value, attention.permute(0, 2, 1))
        out = out.view(m_batchsize, C, 1, height, width)

        out = self.gamma*out + x
        return out
